package RSA;

import java.math.BigInteger;
import java.security.SecureRandom;
import java.util.Base64;
import java.nio.charset.StandardCharsets;

/**
 * 手搓RSA算法（支持文本加密，含中文）
 * 核心依赖：BigInteger（大整数运算）、PKCS1填充（基础版）
 * 生产环境禁用！仅用于学习原理
 */
public class HandwrittenRSA {
    // ===================== 1. 核心常量（RSA规范）=====================
    private static final int KEY_SIZE = 512; // 密钥长度（512位，演示用；实际至少2048位）
    private static final int PADDING_BLOCK_SIZE = 11; // PKCS1填充占用字节数
    private static final BigInteger E = BigInteger.valueOf(65537); // 公钥e（常用值，高效且安全）
    private static final SecureRandom RANDOM = new SecureRandom(); // 安全随机数

    // 密钥对（公钥：e+n；私钥：d+n）
    private BigInteger publicKeyE;
    private BigInteger privateKeyD;
    private BigInteger modulusN;

    // ===================== 2. 核心工具：PKCS1填充/去填充（解决明文长度+安全）=====================
    /**
     * PKCS1填充（基础版）：明文字节 → 填充后字节（用于加密）
     * 规则：0x00 + 0x02 + 随机非零字节 + 0x00 + 原始明文
     */
    private byte[] pkcs1Pad(byte[] data, int maxBlockSize) {
        if (data.length > maxBlockSize - PADDING_BLOCK_SIZE) {
            throw new IllegalArgumentException("明文块过长（最大支持" + (maxBlockSize - PADDING_BLOCK_SIZE) + "字节）");
        }

        byte[] paddedData = new byte[maxBlockSize];
        paddedData[0] = 0x00; // 第1字节固定0x00
        paddedData[1] = 0x02; // 第2字节固定0x02（加密标识）

        // 第3~第10字节：随机非零字节（共8字节）
        int randomLen = maxBlockSize - data.length - 3; // 0x00+0x02+0x00 + 明文
        for (int i = 2; i < 2 + randomLen; i++) {
            byte b;
            do {
                b = (byte) RANDOM.nextInt(256);
            } while (b == 0x00); // 确保随机字节非零
            paddedData[i] = b;
        }

        paddedData[2 + randomLen] = 0x00; // 分隔符：0x00
        System.arraycopy(data, 0, paddedData, 2 + randomLen + 1, data.length); // 复制原始明文

        return paddedData;
    }

    /**
     * PKCS1去填充（基础版）：填充后字节 → 原始明文字节（用于解密）
     */
    /**
     * PKCS1去填充（基础版）：填充后字节 → 原始明文字节（用于解密）
     */
    private byte[] pkcs1Unpad(byte[] paddedData) {
        // 新增：打印填充块前2字节，排查格式问题
        System.out.println("解密后填充块前2字节：0x" + String.format("%02X", paddedData[0]) + " 0x" + String.format("%02X", paddedData[1]));

        // 验证填充格式（第1字节0x00，第2字节0x02）
        if (paddedData[0] != 0x00 || paddedData[1] != 0x02) {
            throw new IllegalArgumentException("无效的填充数据（前2字节应为0x00 0x02，实际为0x" + String.format("%02X", paddedData[0]) + " 0x" + String.format("%02X", paddedData[1]) + "）");
        }

        // 找到分隔符0x00（从第3字节开始找）
        int separatorIndex = -1;
        for (int i = 2; i < paddedData.length; i++) {
            if (paddedData[i] == 0x00) {
                separatorIndex = i;
                break;
            }
        }
        if (separatorIndex == -1) {
            throw new IllegalArgumentException("填充分隔符未找到");
        }

        // 提取原始明文（分隔符后的数据）
        int plainLen = paddedData.length - separatorIndex - 1;
        byte[] plainData = new byte[plainLen];
        System.arraycopy(paddedData, separatorIndex + 1, plainData, 0, plainLen);
        return plainData;
    }

    // ===================== 3. 密钥生成（核心：大质数+模逆）=====================
    /**
     * 生成RSA密钥对（公钥e+n，私钥d+n）
     * 步骤：生成大质数p/q → 计算n=pq → 计算φ(n)=(p-1)(q-1) → 计算d=e^-1 mod φ(n)
     */
    public void generateKeyPair() {
        // 1. 生成两个大质数p和q（512位密钥→p/q各256位伪素数，isProbablePrime(100)表示素性检测可信度100）
        BigInteger p = new BigInteger(KEY_SIZE / 2, 100, RANDOM);
        BigInteger q = new BigInteger(KEY_SIZE / 2, 100, RANDOM);

        // 2. 计算模数n（公钥/私钥共用）和欧拉函数φ(n)
        modulusN = p.multiply(q);
        BigInteger phiN = p.subtract(BigInteger.ONE).multiply(q.subtract(BigInteger.ONE));

        // 3. 计算私钥d（e的模逆元，必须满足e和φ(n)互质）
        privateKeyD = E.modInverse(phiN);

        // 4. 公钥e（固定为65537，无需随机生成）
        publicKeyE = E;

        System.out.println("RSA密钥生成完成：");
        System.out.println("公钥e（Base64）：" + encodeBase64(publicKeyE.toByteArray()));
        System.out.println("私钥d（Base64）：" + encodeBase64(privateKeyD.toByteArray()));
        System.out.println("模数n（Base64）：" + encodeBase64(modulusN.toByteArray()));
        System.out.println("=====================================");
    }

    // ===================== 4. 加密：文本→密文（Base64）=====================
    public String encrypt(String plainText) {
        if (publicKeyE == null || modulusN == null) {
            throw new IllegalStateException("请先生成密钥对");
        }

        // 1. 文本→UTF-8字节
        byte[] plainBytes = plainText.getBytes(StandardCharsets.UTF_8);

        // 2. 计算最大明文块长度（n的字节数 - 11，11是PKCS1填充固定占用）
        int nByteLength = modulusN.toByteArray().length;
        int maxPlainBlockSize = nByteLength - PADDING_BLOCK_SIZE;

        // 3. 分块+填充+加密
        StringBuilder cipherBuilder = new StringBuilder();
        for (int i = 0; i < plainBytes.length; i += maxPlainBlockSize) {
            // 截取当前块明文
            int blockLen = Math.min(maxPlainBlockSize, plainBytes.length - i);
            byte[] plainBlock = new byte[blockLen];
            System.arraycopy(plainBytes, i, plainBlock, 0, blockLen);

            // PKCS1填充
            byte[] paddedBlock = pkcs1Pad(plainBlock, nByteLength);

            // 加密核心：m^e mod n（BigInteger自动处理大整数运算）
            BigInteger m = new BigInteger(1, paddedBlock); // 1表示正数
            BigInteger c = m.modPow(publicKeyE, modulusN);

            // 密文块→Base64（避免字节拼接错乱）
            cipherBuilder.append(encodeBase64(c.toByteArray())).append("|"); // 用|分隔块
        }

        // 去除最后一个|，返回完整密文
        return cipherBuilder.substring(0, cipherBuilder.length() - 1);
    }

    // ===================== 5. 解密：密文（Base64）→文本=====================
    /**
     * 解密：密文（Base64）→文本
     */
    public String decrypt(String cipherText) {
        if (privateKeyD == null || modulusN == null) {
            throw new IllegalStateException("请先生成密钥对");
        }

        // 1. 分割密文块（按|拆分）
        String[] cipherBlocks = cipherText.split("\\|");
        byte[] allPlainBytes = new byte[0];
        // 关键：获取模数n的字节数（加密时的填充块长度，必须一致）
        int nByteLength = modulusN.toByteArray().length;

        // 2. 逐块解密+去填充
        for (String blockStr : cipherBlocks) {
            // Base64解码→密文块字节→BigInteger
            byte[] cipherBlockBytes = decodeBase64(blockStr);
            BigInteger c = new BigInteger(1, cipherBlockBytes);

            // 解密核心：c^d mod n
            BigInteger m = c.modPow(privateKeyD, modulusN);
            byte[] mBytes = m.toByteArray(); // 可能丢失前导0x00，长度可能小于nByteLength

            // ========== 新增：补全前导零，恢复填充块长度 ==========
            byte[] paddedBlock;
            if (mBytes.length == nByteLength) {
                // 长度刚好，直接使用
                paddedBlock = mBytes;
            } else if (mBytes.length < nByteLength) {
                // 长度不足：前导补零（恢复加密时的0x00开头）
                paddedBlock = new byte[nByteLength];
                // 从末尾复制mBytes到paddedBlock，前面补0x00
                System.arraycopy(mBytes, 0, paddedBlock, nByteLength - mBytes.length, mBytes.length);
            } else {
                // 长度超出（极罕见，可能是密钥生成异常），截取后nByteLength字节
                paddedBlock = new byte[nByteLength];
                System.arraycopy(mBytes, mBytes.length - nByteLength, paddedBlock, 0, nByteLength);
            }
            // =====================================================

            // PKCS1去填充（现在paddedBlock格式正确了）
            byte[] plainBlock = pkcs1Unpad(paddedBlock);

            // 拼接明文块
            byte[] temp = new byte[allPlainBytes.length + plainBlock.length];
            System.arraycopy(allPlainBytes, 0, temp, 0, allPlainBytes.length);
            System.arraycopy(plainBlock, 0, temp, allPlainBytes.length, plainBlock.length);
            allPlainBytes = temp;
        }

        // 3. 明文字节→UTF-8文本
        return new String(allPlainBytes, StandardCharsets.UTF_8);
    }

    // ===================== 辅助工具：Base64编码/解码（方便传输存储）=====================
    private String encodeBase64(byte[] data) {
        return Base64.getEncoder().encodeToString(data);
    }

    private byte[] decodeBase64(String dataStr) {
        return Base64.getDecoder().decode(dataStr);
    }

    // ===================== 测试：加密中文文本=====================
    public static void main(String[] args) {
        try {
            // 1. 初始化手搓RSA
            HandwrittenRSA rsa = new HandwrittenRSA();

            // 2. 生成密钥对（512位，演示用）
            rsa.generateKeyPair();

            // 3. 待加密文本（支持中文、特殊字符）
            String plainText = "Hello 手搓RSA！这是一段测试文本：123456 + @#$%^&*()";
            System.out.println("原始明文：" + plainText);
            System.out.println("=====================================");

            // 4. 加密
            String cipherText = rsa.encrypt(plainText);
            System.out.println("加密后的密文（Base64，按块分隔）：" + cipherText);
            System.out.println("=====================================");

            // 5. 解密
            String decryptedText = rsa.decrypt(cipherText);
            System.out.println("解密后的明文：" + decryptedText);
            System.out.println("=====================================");

            // 验证一致性
            if (plainText.equals(decryptedText)) {
                System.out.println("✅ RSA 加密解密成功！明文一致");
            } else {
                System.out.println("❌ RSA 加密解密失败！明文不一致");
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}